{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0,  2,  4,  6,  8],\n",
       "         [10, 12, 14, 16, 18],\n",
       "         [20, 22, 24, 26, 28],\n",
       "         [30, 32, 34, 36, 38],\n",
       "         [40, 42, 44, 46, 48]],\n",
       "\n",
       "        [[50, 52, 54, 56, 58],\n",
       "         [60, 62, 64, 66, 68],\n",
       "         [70, 72, 74, 76, 78],\n",
       "         [80, 82, 84, 86, 88],\n",
       "         [90, 92, 94, 96, 98]]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 2\n",
    "C = 3\n",
    "H = 5\n",
    "W = 5\n",
    "x = torch.arange(N*H*W).reshape(N, H, W)*2\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[20,  6, 11, 14, 11, 12, 15, 19],\n",
       "         [ 4,  6,  8, 14, 15,  3, 16, 18]]), torch.Size([2, 8]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indices = torch.randint(H * W, (N, 8), dtype=torch.long)\n",
    "indices, indices.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[40, 12, 22, 28, 22, 24, 30, 38],\n",
       "         [58, 62, 66, 78, 80, 56, 82, 86]]), torch.Size([2, 8]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga = torch.gather(x.view(N, H * W), 1, indices).view(N, -1)\n",
    "ga, ga.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0]]), tensor([[[ 0,  2,  4,  6,  8],\n",
       "          [10, 12, 14, 16, 18],\n",
       "          [20, 22, 24, 26, 28],\n",
       "          [30, 32, 34, 36, 38],\n",
       "          [40, 42, 44, 46, 48]],\n",
       " \n",
       "         [[50, 52, 54, 56, 58],\n",
       "          [60, 62, 64, 66, 68],\n",
       "          [70, 72, 74, 76, 78],\n",
       "          [80, 82, 84, 86, 88],\n",
       "          [90, 92, 94, 96, 98]]]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga = ga / 100\n",
    "ga, x "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0,  2,  4,  6,  8],\n",
       "         [10, 12, 14, 16, 18],\n",
       "         [20, 22, 24, 26, 28],\n",
       "         [30, 32, 34, 36, 38],\n",
       "         [40, 42, 44, 46, 48]]), tensor([[50, 52, 54, 56, 58],\n",
       "         [60, 62, 64, 66, 68],\n",
       "         [70, 72, 74, 76, 78],\n",
       "         [80, 82, 84, 86, 88],\n",
       "         [90, 92, 94, 96, 98]]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x2 = torch.arange(N*C*H*W).reshape(N, C, H, W)*2\n",
    "x2[0][0], x2[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 8])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indices.unsqueeze(1).expand(-1, C, -1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 40,  12,  22,  28,  22,  24,  30,  38],\n",
       "          [ 90,  62,  72,  78,  72,  74,  80,  88],\n",
       "          [140, 112, 122, 128, 122, 124, 130, 138]],\n",
       " \n",
       "         [[158, 162, 166, 178, 180, 156, 182, 186],\n",
       "          [208, 212, 216, 228, 230, 206, 232, 236],\n",
       "          [258, 262, 266, 278, 280, 256, 282, 286]]]), torch.Size([2, 3, 8]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga = torch.gather(x2.view(N, C, -1), 2, indices.unsqueeze(1).expand(-1, C, -1))\n",
    "ga, ga.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 40,  12,  22,  28,  22,  24,  30,  38],\n",
       "         [158, 162, 166, 178, 180, 156, 182, 186]]),\n",
       " tensor([[ 90,  62,  72,  78,  72,  74,  80,  88],\n",
       "         [208, 212, 216, 228, 230, 206, 232, 236]]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga[:, 0, :], ga[:, 1, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gather Backward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "x3 = x2.float().requires_grad_(True)\n",
    "ga = torch.gather(x3.view(N, C, -1), 2, indices.unsqueeze(1).expand(-1, C, -1))\n",
    "x3.sum().backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 40.,  12.,  22.,  28.,  22.,  24.,  30.,  38.],\n",
       "         [ 90.,  62.,  72.,  78.,  72.,  74.,  80.,  88.],\n",
       "         [140., 112., 122., 128., 122., 124., 130., 138.]],\n",
       "\n",
       "        [[158., 162., 166., 178., 180., 156., 182., 186.],\n",
       "         [208., 212., 216., 228., 230., 206., 232., 236.],\n",
       "         [258., 262., 266., 278., 280., 256., 282., 286.]]],\n",
       "       grad_fn=<GatherBackward>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[16,  5,  6,  8, 13,  3, 22, 16, 24, 16, 20, 18, 22, 11,  2,  3,  2, 12,\n",
       "           4,  8, 23, 15, 12,  3,  2, 13, 10,  7,  0, 11],\n",
       "         [ 5, 13,  7, 20,  6,  1, 19, 18, 17,  8, 11,  3, 10,  7,  9, 19, 19, 11,\n",
       "          24,  7, 21, 10,  7,  1, 12, 21,  0, 17,  8, 12]]),\n",
       " tensor([[32, 10, 12, 16, 26,  6, 44, 32, 48, 32, 40, 36, 44, 22,  4,  6,  4, 24,\n",
       "           8, 16, 46, 30, 24,  6,  4, 26, 20, 14,  0, 22],\n",
       "         [60, 76, 64, 90, 62, 52, 88, 86, 84, 66, 72, 56, 70, 64, 68, 88, 88, 72,\n",
       "          98, 64, 92, 70, 64, 52, 74, 92, 50, 84, 66, 74]]))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 2\n",
    "C = 3\n",
    "H = 5\n",
    "W = 5\n",
    "x = torch.arange(N*H*W).reshape(N, H, W)*2\n",
    "indices = torch.randint(H * W, (N, 30), dtype=torch.long)\n",
    "ga = torch.gather(x.view(N, H * W), 1, indices).view(N, -1)\n",
    "indices, ga"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value ; tensor([[48, 46, 44, 44, 40],\n",
      "        [98, 92, 92, 90, 88]]) torch.Size([2, 5])\n",
      "Idx ; tensor([[ 8, 20,  6, 12, 10],\n",
      "        [18, 25, 20,  3, 15]]) torch.Size([2, 5])\n"
     ]
    }
   ],
   "source": [
    "v, i = ga.topk(5)\n",
    "print(\"Value ;\", v, v.shape)\n",
    "print(\"Idx ;\",i, i.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[24, 23, 22, 22, 20],\n",
       "        [24, 21, 21, 20, 19]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reverse_indices = torch.gather(indices, 1, i)\n",
    "reverse_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[48, 46, 44, 44, 40],\n",
       "        [98, 92, 92, 90, 88]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.gather(x.view(N, H * W), 1, reverse_indices).view(N, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 0,  1,  2,  3],\n",
       "          [ 4,  5,  6,  7],\n",
       "          [ 8,  9, 10, 11],\n",
       "          [12, 13, 14, 15]],\n",
       " \n",
       "         [[16, 17, 18, 19],\n",
       "          [20, 21, 22, 23],\n",
       "          [24, 25, 26, 27],\n",
       "          [28, 29, 30, 31]]]),\n",
       " tensor([[[  0,   1,   2,   3,   4,   5,   6,   7],\n",
       "          [  8,   9,  10,  11,  12,  13,  14,  15],\n",
       "          [ 16,  17,  18,  19,  20,  21,  22,  23],\n",
       "          [ 24,  25,  26,  27,  28,  29,  30,  31],\n",
       "          [ 32,  33,  34,  35,  36,  37,  38,  39],\n",
       "          [ 40,  41,  42,  43,  44,  45,  46,  47],\n",
       "          [ 48,  49,  50,  51,  52,  53,  54,  55],\n",
       "          [ 56,  57,  58,  59,  60,  61,  62,  63]],\n",
       " \n",
       "         [[ 64,  65,  66,  67,  68,  69,  70,  71],\n",
       "          [ 72,  73,  74,  75,  76,  77,  78,  79],\n",
       "          [ 80,  81,  82,  83,  84,  85,  86,  87],\n",
       "          [ 88,  89,  90,  91,  92,  93,  94,  95],\n",
       "          [ 96,  97,  98,  99, 100, 101, 102, 103],\n",
       "          [104, 105, 106, 107, 108, 109, 110, 111],\n",
       "          [112, 113, 114, 115, 116, 117, 118, 119],\n",
       "          [120, 121, 122, 123, 124, 125, 126, 127]]]),\n",
       " tensor([[10,  6,  5,  8,  1],\n",
       "         [10,  2,  9,  9, 10]]))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 2\n",
    "C = 1\n",
    "H = 4\n",
    "W = 4\n",
    "x = torch.arange(N*H*W).reshape(N, H, W)\n",
    "indices = torch.randint(H * W, (N, 5), dtype=torch.long)\n",
    "\n",
    "s = [2, 2]\n",
    "x2 = torch.arange(N*H*W*s[0]*s[1]).reshape(N, H*s[0], W*s[1])\n",
    "x, x2, indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[10,  6,  5,  8,  1],\n",
      "        [10,  2,  9,  9, 10]]) \n",
      " tensor([[36, 20, 18, 32,  2],\n",
      "        [36,  4, 34, 34, 36]])\n"
     ]
    }
   ],
   "source": [
    "ys = (indices // W) * s[0]\n",
    "xs = (indices % W) * s[1]\n",
    "\n",
    "W *= s[1]\n",
    "\n",
    "new_ind = (ys * W) + xs\n",
    "print(indices, \"\\n\", new_ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Infer, Issues #5\n",
    "import torch\n",
    "N = 2\n",
    "C = 3\n",
    "H = 4\n",
    "W = 4\n",
    "x = torch.arange(N*C*H*W).reshape(N, C, H, W).float()\n",
    "indices = torch.randint(H * W, (N, 5), dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ -6., -15., -12., -11.,  -2.],\n",
       "         [-22., -31., -28., -27., -18.],\n",
       "         [-38., -47., -44., -43., -34.]],\n",
       "\n",
       "        [[-57., -58., -49., -50., -57.],\n",
       "         [-73., -74., -65., -66., -73.],\n",
       "         [-89., -90., -81., -82., -89.]]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ga = torch.gather(x.view(N, C, -1), 2,\n",
    "                  indices.unsqueeze(1).expand(-1, C, -1)) * -1\n",
    "ga"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[  0.,   1.,  -2.,   3.,   4.,   5.,  -6.,   7.,   8.,   9.,  10.,\n",
       "          -11., -12.,  13.,  14., -15.],\n",
       "         [ 16.,  17., -18.,  19.,  20.,  21., -22.,  23.,  24.,  25.,  26.,\n",
       "          -27., -28.,  29.,  30., -31.],\n",
       "         [ 32.,  33., -34.,  35.,  36.,  37., -38.,  39.,  40.,  41.,  42.,\n",
       "          -43., -44.,  45.,  46., -47.]],\n",
       "\n",
       "        [[ 48., -49., -50.,  51.,  52.,  53.,  54.,  55.,  56., -57., -58.,\n",
       "           59.,  60.,  61.,  62.,  63.],\n",
       "         [ 64., -65., -66.,  67.,  68.,  69.,  70.,  71.,  72., -73., -74.,\n",
       "           75.,  76.,  77.,  78.,  79.],\n",
       "         [ 80., -81., -82.,  83.,  84.,  85.,  86.,  87.,  88., -89., -90.,\n",
       "           91.,  92.,  93.,  94.,  95.]]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rend_pred = x.view(N, C, -1).scatter_(2, indices.unsqueeze(1).expand(-1, C, -1), ga)\n",
    "# rend_pred = rend_pred.reshape((N, C, H, W))\n",
    "rend_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
